import easyocr
import cv2
import matplotlib.pyplot as plt

# 1. 初始化Reader（首次运行会自动下载模型）
reader = easyocr.Reader(
    ['ch_sim', 'en'],  # 中文简体和英文
    gpu=False,          # 启用GPU可加速（需安装CUDA）
    model_storage_directory=r'D:\ideaSpace\MyPython\models\easyocr_models'  # 自定义模型路径
)

# 2. 执行OCR并可视化
def easyocr_recognize(image_path):
    # 读取图像
    image = cv2.imread(image_path)

    # 执行OCR
    results = reader.readtext(
        image,
        decoder='beamsearch',  # 提高精度（可选：'greedy'更快）
        batch_size=4,          # 并行处理
        detail=1               # 返回详细信息
    )

    # 绘制结果
    for (bbox, text, prob) in results:
        (tl, tr, br, bl) = bbox
        tl = (int(tl[0]), int(tl[1]))
        br = (int(br[0]), int(br[1]))
        cv2.rectangle(image, tl, br, (0, 255, 0), 2)
        cv2.putText(image, f"{text} ({prob:.2f})", (tl[0], tl[1] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)

    # 显示
    plt.figure(figsize=(10, 10))
    plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    plt.axis('off')
    plt.show()

    # 返回结构化结果
    return [(text, prob, bbox) for (bbox, text, prob) in results]

# 3. 执行识别
image_path = r"D:\ideaSpace\rag-in-action-master\wukong.jpg"
results = easyocr_recognize(image_path)
print("识别结果（文本, 置信度, 位置）：")
for i, (text, prob, bbox) in enumerate(results, 1):
    print(f"{i}. {text} (置信度: {prob:.2f})")